from langchain_community.document_loaders import TextLoader
from langchain.text_splitter import CharacterTextSplitter


import abc

# 文件工具抽象类
class ABCFile(metaclass=abc.ABCMeta):
    
    @abc.abstractmethod
    def fileLoader(self):
        pass 

    # 自动分割文档
    @abc.abstractmethod
    def fileSplitAuto(self):
        pass     
    
    # 按照目录分割文档
    @abc.abstractmethod
    def fileSplitMenu(self):
        pass     
    
    # 使用分隔符分割文档
    @abc.abstractmethod
    def fileSplitChar(self):
        pass 
    
class FileUtil():
    
    def doc2List(self,docs):
        ret = []
        for doc in docs:
            ret.append(doc.page_content)
        return ret

# txt文件工具
class TxtUtil(FileUtil,ABCFile):
    
    def __init__(self) -> None:
        pass
    
    def fileLoader(self,file_path):
        # 实例化TextLoader对象
        loader = TextLoader(file_path,encoding="utf-8")
        # 加载文档
        docs = loader.load()   
        
        return docs    
    
    def fileSplitAuto(self,docs):
        # 初始化分割器
        # split = CharacterTextSplitter(splitchar,chunk_size=chunk_size, chunk_overlap=chunk_overlap)
        # chunks = split.split_documents(docs)
        chunks = None
        return chunks   
    
    def fileSplitMenu(self,docs):
        # 初始化分割器
        # split = CharacterTextSplitter(splitchar,chunk_size=chunk_size, chunk_overlap=chunk_overlap)
        # chunks = split.split_documents(docs)
        chunks = None
        return chunks   
    
    def fileSplitChar(self,docs,splitchar,chunk_size,chunk_overlap):
        # 初始化分割器
        split = CharacterTextSplitter(splitchar,chunk_size=chunk_size, chunk_overlap=chunk_overlap)
        chunks = split.split_documents(docs)
        
        return chunks   
    

# Doc文件工具
class DocUtil(FileUtil):
    
    def __init__(self) -> None:
        pass
    
    def fileLoader(self,file_path):
        # 实例化TextLoader对象
        loader = TextLoader(file_path,encoding="utf-8")
        # 加载文档
        docs = loader.load()   
        
        return docs    
    
    def fileSplitAuto(self,docs):
        # 初始化分割器
        split = CharacterTextSplitter(splitchar,chunk_size=chunk_size, chunk_overlap=chunk_overlap)
        chunks = split.split_documents(docs)
        
        return chunks   
    
    def fileSplitMenu(self,docs):
        # 初始化分割器
        # split = CharacterTextSplitter(splitchar,chunk_size=chunk_size, chunk_overlap=chunk_overlap)
        # chunks = split.split_documents(docs)
        chunks = None
        return chunks   
    
    def fileSplitChar(self,docs,splitchar,chunk_size,chunk_overlap):
        # 初始化分割器
        split = CharacterTextSplitter(splitchar,chunk_size=chunk_size, chunk_overlap=chunk_overlap)
        chunks = split.split_documents(docs)
        
        return chunks 


# Doc文件工具
class DocxUtil(FileUtil):
    
    def __init__(self) -> None:
        pass
    
    def fileLoader(self,file_path):
        # 实例化TextLoader对象
        loader = TextLoader(file_path,encoding="utf-8")
        # 加载文档
        docs = loader.load()   
        
        return docs    
    
    def fileSplitAuto(self,docs):
        # 初始化分割器
        split = CharacterTextSplitter(splitchar,chunk_size=chunk_size, chunk_overlap=chunk_overlap)
        chunks = split.split_documents(docs)
        
        return chunks   
    
    def fileSplitMenu(self,docs):
        # 初始化分割器
        # split = CharacterTextSplitter(splitchar,chunk_size=chunk_size, chunk_overlap=chunk_overlap)
        # chunks = split.split_documents(docs)
        chunks = None
        return chunks   
    
    def fileSplitChar(self,docs,splitchar,chunk_size,chunk_overlap):
        # 初始化分割器
        split = CharacterTextSplitter(splitchar,chunk_size=chunk_size, chunk_overlap=chunk_overlap)
        chunks = split.split_documents(docs)
        
        return chunks 

# Doc文件工具
class PdfUtil(FileUtil):
    
    def __init__(self) -> None:
        pass
    
    def fileLoader(self,file_path):
        # 实例化TextLoader对象
        loader = TextLoader(file_path,encoding="utf-8")
        # 加载文档
        docs = loader.load()   
        
        return docs    
    
    def fileSplitAuto(self,docs):
        # 初始化分割器
        split = CharacterTextSplitter(splitchar,chunk_size=chunk_size, chunk_overlap=chunk_overlap)
        chunks = split.split_documents(docs)
        
        return chunks   
    
    def fileSplitMenu(self,docs):
        # 初始化分割器
        # split = CharacterTextSplitter(splitchar,chunk_size=chunk_size, chunk_overlap=chunk_overlap)
        # chunks = split.split_documents(docs)
        chunks = None
        return chunks   
    
    def fileSplitChar(self,docs,splitchar,chunk_size,chunk_overlap):
        # 初始化分割器
        split = CharacterTextSplitter(splitchar,chunk_size=chunk_size, chunk_overlap=chunk_overlap)
        chunks = split.split_documents(docs)
        
        return chunks 
 
class FileFactory:
    def get_sclass(self,type):
        if type == "txt":
            return TxtUtil()
        elif type == "doc":
            return DocUtil()
        
        elif type == "docx":
            return DocxUtil()
        
        elif type == "pdf":
            return PdfUtil()